# modules/auth.py
import streamlit as st
import sqlite3
import hashlib
import os

# 数据库路径配置
DB_PATH = "data/users.db"


def create_users_table():
    """创建用户数据表"""
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS users
                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
                  username TEXT UNIQUE NOT NULL,
                  password_hash TEXT NOT NULL,
                  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
    conn.commit()
    conn.close()


def hash_password(password):
    """使用SHA-256加盐哈希加密密码"""
    salt = os.urandom(16)  # 生成随机盐值
    pwd_hash = hashlib.pbkdf2_hmac('sha256',
                                   password.encode('utf-8'),
                                   salt,
                                   100000)
    return salt.hex() + pwd_hash.hex()


def verify_password(stored_hash, input_password):
    """验证密码哈希"""
    try:
        salt = bytes.fromhex(stored_hash[:32])  # 提取盐值（前32字符）
        input_hash = hashlib.pbkdf2_hmac('sha256',
                                         input_password.encode('utf-8'),
                                         salt,
                                         100000).hex()
        return input_hash == stored_hash[32:]
    except:
        return False


def register_user():
    """用户注册对话框"""
    with st.expander("新用户注册", expanded=True):
        with st.form("注册表单"):
            new_username = st.text_input("用户名（需唯一）")
            new_password = st.text_input("密码", type="password")
            submit = st.form_submit_button("提交注册")

            if submit:
                if not new_username or not new_password:
                    st.error("用户名和密码不能为空")
                    return

                try:
                    conn = sqlite3.connect(DB_PATH)
                    c = conn.cursor()

                    # 检查用户名是否已存在
                    c.execute("SELECT 1 FROM users WHERE username=?", (new_username,))
                    if c.fetchone():
                        st.error("用户名已存在")
                        return

                    # 插入新用户
                    pwd_hash = hash_password(new_password)
                    c.execute("INSERT INTO users (username, password_hash) VALUES (?, ?)",
                              (new_username, pwd_hash))
                    conn.commit()
                    st.success("注册成功！请返回登录")

                except Exception as e:
                    st.error(f"注册失败: {str(e)}")
                finally:
                    conn.close()


def authenticate_user(username, password):
    """用户认证"""
    try:
        conn = sqlite3.connect(DB_PATH)
        c = conn.cursor()
        c.execute("SELECT password_hash FROM users WHERE username=?", (username,))
        result = c.fetchone()

        if not result:
            return False

        stored_hash = result[0]
        return verify_password(stored_hash, password)

    except Exception as e:
        st.error(f"认证错误: {str(e)}")
        return False
    finally:
        conn.close()